#pragma once
#include <memory>
#include <filesystem>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <rclcpp/rclcpp.hpp>
#include <opencv2/opencv.hpp>
#include <nav_msgs/msg/occupancy_grid.hpp>
#include <sensor_msgs/msg/point_cloud2.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/radius_outlier_removal.h>

using point_cloud_xyz = pcl::PointCloud<pcl::PointXYZ>;

std::shared_ptr<nav_msgs::msg::OccupancyGrid> load_grid_from_image(
    std::string header_frame_id,
    std::string map_file_path,
    double resolution = 0.1
) {
    cv::Mat image = cv::imread(map_file_path, cv::IMREAD_GRAYSCALE);
    if (image.empty()) {
        std::cout << "Failed to load map image: " << map_file_path << std::endl;
        return nullptr;
    }

    std::filesystem::path p(map_file_path);
    std::string filename = p.stem().string();
    size_t last_underscore = filename.rfind('_');
    size_t second_last_underscore = filename.rfind('_', last_underscore - 1);
    std::string num1 = filename.substr(second_last_underscore + 1, last_underscore - second_last_underscore - 1);
    std::string num2 = filename.substr(last_underscore + 1);
    double px_min = std::stod(num1) / 100.0;
    double py_min = std::stod(num2) / 100.0;

    std::cout << "Loaded map image from: " << map_file_path << std::endl;

    int width = image.cols;
    int height = image.rows;
    auto grid = std::make_shared<nav_msgs::msg::OccupancyGrid>();
    grid->header.frame_id = header_frame_id;
    grid->header.stamp = rclcpp::Time(0);
    grid->info.resolution = resolution; // 10 cm resolution
    grid->info.width = width;
    grid->info.height = height;
    grid->info.origin.position.x = px_min; // Adjust as needed
    grid->info.origin.position.y = py_min; // Adjust as needed
    grid->info.origin.position.z = 0.0;
    grid->data.resize(width * height);

    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            int index = (height - 1 - y) * width + x;
            uchar pixel_value = image.at<uchar>(y, x);
            if (pixel_value == 255) {
                grid->data[index] = 100; // Occupied
            } else {
                grid->data[index] = 0; // Unoccupied
            }
        }
    }
    return grid;
}

point_cloud_xyz::Ptr clip_cloud_and_reverse(
    point_cloud_xyz::Ptr cloud_ptr,
    double min_pz,
    double max_pz,
    bool reverse = false
) {
    /* Clip the cloud_ptr_ based on the min and max pz values obtained from the path file. */
    
    auto filtered_cloud = std::make_shared<point_cloud_xyz>();
    filtered_cloud->header = cloud_ptr->header;
    filtered_cloud->is_dense = true;
    for (auto &point : cloud_ptr->points) {
        // Clip based on min and max pz with gain
        if (point.z < min_pz || point.z > max_pz) {
            continue;
        }
        if (reverse) {
            // Reverse the point coordinates
            point.x = -point.x;
            // point.y = -point.y;
            point.z = -point.z;
        }
        filtered_cloud->points.push_back(point);
    }
    filtered_cloud->width = filtered_cloud->points.size();
    filtered_cloud->height = 1;
    filtered_cloud->is_dense = true;
    cloud_ptr = filtered_cloud;

    return cloud_ptr;
}

std::shared_ptr<nav_msgs::msg::OccupancyGrid> build_2d_map(
    point_cloud_xyz::Ptr cloud_ptr,
    double resolution,
    double grid_value_threshold
) {
    // Setup the occupancy grid parameters
    float min_x = 1e9, max_x = -1e9, min_y = 1e9, max_y = -1e9;
    for (const auto &point : cloud_ptr->points) {
        min_x = std::min(min_x, point.x);
        max_x = std::max(max_x, point.x);
        min_y = std::min(min_y, point.y);
        max_y = std::max(max_y, point.y);
    }
    int width = static_cast<int>((max_x - min_x) / resolution) + 1;
    int height = static_cast<int>((max_y - min_y) / resolution) + 1;

    std::vector<int> cummulateted_points(width * height, 0);
    std::vector<std::vector<int>> hist(height, std::vector<int>(width, 0));
    // Count points in each grid cell
    for (const auto &point : cloud_ptr->points) {
        int grid_x = static_cast<int>((point.x - min_x) / resolution);
        int grid_y = static_cast<int>((point.y - min_y) / resolution);
        if (grid_x >= 0 && grid_x < width && grid_y >= 0 && grid_y < height) {
            ++hist[grid_y][grid_x];
        }
    }
    // Compute square closure
    int available_points = 0;
    int min_x_idx = width, min_y_idx = height, max_x_idx = 0, max_y_idx = 0;
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            if (hist[y][x] > grid_value_threshold) {
                min_x_idx = std::min(min_x_idx, x);
                min_y_idx = std::min(min_y_idx, y);
                max_x_idx = std::max(max_x_idx, x);
                max_y_idx = std::max(max_y_idx, y);
                ++available_points;
            }
        }
    }
    if (available_points == 0) {
        return nullptr; // No points above the threshold
    }
    int padding = 1;  // Padding grid cells
    min_x_idx = std::max(0, min_x_idx - padding);
    min_y_idx = std::max(0, min_y_idx - padding);
    max_x_idx = std::min(width - 1, max_x_idx + padding);
    max_y_idx = std::min(height - 1, max_y_idx + padding);
    
    // Convert the histogram to an occupancy grid
    int new_width = max_x_idx - min_x_idx + 1;
    int new_height = max_y_idx - min_y_idx + 1;
    auto grid = std::make_shared<nav_msgs::msg::OccupancyGrid>();
    grid->header.frame_id = "map";
    grid->header.stamp = rclcpp::Time(0);
    grid->info.resolution = resolution;
    grid->info.width = new_width;
    grid->info.height = new_height;
    grid->info.origin.position.x = min_x + min_x_idx * resolution;
    grid->info.origin.position.y = min_y + min_y_idx * resolution;
    grid->data.assign(new_width * new_height, 0);

    for (int y = min_y_idx; y <= max_y_idx; ++y) {
        for (int x = min_x_idx; x <= max_x_idx; ++x) {
            int value = (hist[y][x] > grid_value_threshold) ? 100 : 0;
            int new_idx = (y - min_y_idx) * new_width + (x - min_x_idx);
            grid->data[new_idx] = value;
        }
    }

    return grid;
}

point_cloud_xyz::Ptr apply_voxel_filter(
    point_cloud_xyz::Ptr& cloud_ptr,
    double leaf_size_x,
    double leaf_size_y,
    double leaf_size_z
) {
    pcl::VoxelGrid<pcl::PointXYZ> voxel_filter;
    voxel_filter.setInputCloud(cloud_ptr);
    voxel_filter.setLeafSize(leaf_size_x, leaf_size_y, leaf_size_z);
    point_cloud_xyz::Ptr filtered_cloud(new point_cloud_xyz);
    voxel_filter.filter(*filtered_cloud);
    return filtered_cloud;
}

point_cloud_xyz::Ptr apply_voxel_filter(point_cloud_xyz::Ptr& cloud_ptr, double leaf_size = 0.1) {
    return apply_voxel_filter(cloud_ptr, leaf_size, leaf_size, leaf_size);
}

point_cloud_xyz::Ptr apply_statistical_outlier_removal(
    point_cloud_xyz::Ptr& cloud_ptr,
    int mean_k = 50,
    double std_dev_mul_thresh = 1.0
) {
    pcl::StatisticalOutlierRemoval<pcl::PointXYZ> sor;
    sor.setInputCloud(cloud_ptr);
    sor.setMeanK(mean_k);
    sor.setStddevMulThresh(std_dev_mul_thresh);
    point_cloud_xyz::Ptr filtered_cloud(new point_cloud_xyz);
    sor.filter(*filtered_cloud);
    return filtered_cloud;
}

point_cloud_xyz::Ptr apply_radius_outlier_removal(
    point_cloud_xyz::Ptr& cloud_ptr,
    double radius = 0.1,
    int min_neighbors = 2
) {
    pcl::RadiusOutlierRemoval<pcl::PointXYZ> outrem;
    outrem.setInputCloud(cloud_ptr);
    outrem.setRadiusSearch(radius);
    outrem.setMinNeighborsInRadius(min_neighbors);
    point_cloud_xyz::Ptr filtered_cloud(new point_cloud_xyz);
    outrem.filter(*filtered_cloud);
    return filtered_cloud;
}
